# -*- coding: utf-8 -*-
"""
安全中间件

提供Web应用的安全保护功能，包括：
- 安全响应头设置
- 请求限流和防护
- XSS和CSRF防护
- 内容类型验证
- 请求大小限制
"""

import time
import hashlib
from typing import Callable, Dict, Set
from fastapi import Request, Response, HTTPException, status
from starlette.middleware.base import BaseHTTPMiddleware
from collections import defaultdict, deque
import logging

logger = logging.getLogger(__name__)


class SecurityMiddleware(BaseHTTPMiddleware):
    """安全中间件"""

    def __init__(
            self,
            app,
            rate_limit_requests: int = 100,  # 每分钟请求数限制
            rate_limit_window: int = 60,  # 限流窗口（秒）
            max_request_size: int = 10 * 1024 * 1024,  # 最大请求大小（10MB）
            enable_security_headers: bool = True,
            enable_rate_limiting: bool = True,
            enable_request_validation: bool = True
    ):
        super().__init__(app)
        self.rate_limit_requests = rate_limit_requests
        self.rate_limit_window = rate_limit_window
        self.max_request_size = max_request_size
        self.enable_security_headers = enable_security_headers
        self.enable_rate_limiting = enable_rate_limiting
        self.enable_request_validation = enable_request_validation

        # 限流存储（实际应用中应使用Redis）
        self.rate_limit_data: Dict[str, deque] = defaultdict(deque)

        # 恶意请求检测
        self.suspicious_patterns = {
            "sql_injection": [
                "union select", "drop table", "delete from",
                "insert into", "update set", "create table"
            ],
            "xss": [
                "<script", "javascript:", "onclick=", "onerror=",
                "onload=", "eval(", "alert("
            ],
            "path_traversal": [
                "../", "..\\", "/etc/passwd", "/etc/shadow",
                "boot.ini", "win.ini"
            ]
        }

    async def dispatch(self, request: Request, call_next: Callable) -> Response:
        """处理安全检查"""

        # 获取客户端IP
        client_ip = self._get_client_ip(request)

        try:
            # 1. 请求大小验证
            if self.enable_request_validation:
                await self._validate_request_size(request)

            # 2. 恶意请求检测
            if self.enable_request_validation:
                await self._detect_malicious_requests(request)

            # 3. 限流检查
            if self.enable_rate_limiting:
                await self._check_rate_limit(client_ip)

            # 处理请求
            response = await call_next(request)

            # 4. 设置安全响应头
            if self.enable_security_headers:
                self._set_security_headers(response)

            return response

        except HTTPException:
            # 记录安全事件
            logger.warning(f"安全检查失败: {client_ip} - {request.method} {request.url.path}")
            raise
        except Exception as e:
            logger.error(f"安全中间件异常: {e}", exc_info=True)
            raise HTTPException(
                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
                detail="安全检查失败"
            )

    def _get_client_ip(self, request: Request) -> str:
        """获取客户端IP地址"""
        if "X-Real-IP" in request.headers:
            return request.headers["X-Real-IP"]
        elif "X-Forwarded-For" in request.headers:
            return request.headers["X-Forwarded-For"].split(",")[0].strip()
        else:
            return request.client.host if request.client else "unknown"

    async def _validate_request_size(self, request: Request):
        """验证请求大小"""
        content_length = request.headers.get("Content-Length")

        if content_length:
            try:
                size = int(content_length)
                if size > self.max_request_size:
                    raise HTTPException(
                        status_code=status.HTTP_413_REQUEST_ENTITY_TOO_LARGE,
                        detail=f"请求大小超过限制: {self.max_request_size} bytes"
                    )
            except ValueError:
                raise HTTPException(
                    status_code=status.HTTP_400_BAD_REQUEST,
                    detail="无效的Content-Length头"
                )

    async def _detect_malicious_requests(self, request: Request):
        """检测恶意请求"""

        # 检查URL路径
        path = request.url.path.lower()
        query = str(request.url.query).lower()

        # 检查路径遍历攻击
        for pattern in self.suspicious_patterns["path_traversal"]:
            if pattern in path or pattern in query:
                logger.warning(f"检测到路径遍历攻击: {request.url}")
                raise HTTPException(
                    status_code=status.HTTP_400_BAD_REQUEST,
                    detail="非法请求路径"
                )

        # 检查XSS攻击
        for pattern in self.suspicious_patterns["xss"]:
            if pattern in query:
                logger.warning(f"检测到XSS攻击: {request.url}")
                raise HTTPException(
                    status_code=status.HTTP_400_BAD_REQUEST,
                    detail="非法请求参数"
                )

        # 检查请求体中的恶意内容
        if request.method in ["POST", "PUT", "PATCH"]:
            await self._check_request_body(request)

    async def _check_request_body(self, request: Request):
        """检查请求体中的恶意内容"""
        try:
            # 如果是文件上传，跳过检查
            content_type = request.headers.get("content-type", "").lower()
            if "multipart/form-data" in content_type or "application/octet-stream" in content_type:
                return
                
            # 使用缓存的请求体（由RequestBodyCacheMiddleware提供）
            from .request_body_cache_middleware import setup_request_body_cache
            body = setup_request_body_cache(request)
            
            if not body:
                # 没有请求体或未缓存，跳过检查
                return
            
            try:
                body_str = body.decode("utf-8", errors="ignore").lower()
            except Exception:
                # 解码失败，可能是二进制数据，跳过检查
                return
            
            # 检查SQL注入
            for pattern in self.suspicious_patterns["sql_injection"]:
                if pattern in body_str:
                    logger.warning(f"检测到SQL注入攻击: {request.url}")
                    raise HTTPException(
                        status_code=status.HTTP_400_BAD_REQUEST,
                        detail="非法请求内容"
                    )
            
            # 检查XSS攻击
            for pattern in self.suspicious_patterns["xss"]:
                if pattern in body_str:
                    logger.warning(f"检测到XSS攻击: {request.url}")
                    raise HTTPException(
                        status_code=status.HTTP_400_BAD_REQUEST,
                        detail="非法请求内容"
                    )
            
        except HTTPException:
            # 重新抛出HTTP异常
            raise
        except Exception as e:
            logger.error(f"请求体检查失败: {e}")

    async def _check_rate_limit(self, client_ip: str):
        """检查限流"""
        current_time = time.time()

        # 清理过期的请求记录
        client_requests = self.rate_limit_data[client_ip]
        while client_requests and current_time - client_requests[0] > self.rate_limit_window:
            client_requests.popleft()

        # 检查请求数量
        if len(client_requests) >= self.rate_limit_requests:
            logger.warning(f"限流触发: {client_ip} - {len(client_requests)} requests")
            raise HTTPException(
                status_code=status.HTTP_429_TOO_MANY_REQUESTS,
                detail=f"请求过于频繁，请在{self.rate_limit_window}秒后重试",
                headers={"Retry-After": str(self.rate_limit_window)}
            )

        # 记录当前请求
        client_requests.append(current_time)

    def _set_security_headers(self, response: Response):
        """设置安全响应头"""

        # 防止MIME类型嗅探
        response.headers["X-Content-Type-Options"] = "nosniff"

        # 防止点击劫持
        response.headers["X-Frame-Options"] = "DENY"

        # XSS保护
        response.headers["X-XSS-Protection"] = "1; mode=block"

        # 引用者策略
        response.headers["Referrer-Policy"] = "strict-origin-when-cross-origin"

        # 内容安全策略（根据实际需求调整）
        response.headers["Content-Security-Policy"] = (
            "default-src 'self'; "
            "script-src 'self' 'unsafe-inline' 'unsafe-eval'; "
            "style-src 'self' 'unsafe-inline'; "
            "img-src 'self' data: https:; "
            "font-src 'self'; "
            "connect-src 'self'; "
            "media-src 'self'; "
            "object-src 'none'; "
            "child-src 'none'; "
            "worker-src 'none'; "
            "frame-ancestors 'none'; "
            "form-action 'self'; "
            "base-uri 'self'; "
            "manifest-src 'self'"
        )

        # 权限策略
        response.headers["Permissions-Policy"] = (
            "camera=(), "
            "microphone=(), "
            "geolocation=(), "
            "payment=(), "
            "usb=(), "
            "magnetometer=(), "
            "gyroscope=(), "
            "accelerometer=()"
        )

        # HSTS（仅在HTTPS下有效）
        response.headers["Strict-Transport-Security"] = (
            "max-age=31536000; includeSubDomains; preload"
        )

        # 移除服务器信息
        if "Server" in response.headers:
            del response.headers["Server"]

    def get_rate_limit_info(self, client_ip: str) -> Dict[str, int]:
        """获取限流信息"""
        current_time = time.time()
        client_requests = self.rate_limit_data[client_ip]

        # 清理过期请求
        while client_requests and current_time - client_requests[0] > self.rate_limit_window:
            client_requests.popleft()

        return {
            "requests_made": len(client_requests),
            "requests_limit": self.rate_limit_requests,
            "reset_time": int(current_time + self.rate_limit_window),
            "remaining": max(0, self.rate_limit_requests - len(client_requests))
        }
